# import numpy as np
# import matplotlib.pyplot as plt
# from TDMA import run_tdma_simulation
# from env import ENVIRONMENT
#
# def calculate_TONT(env):
#     p = 2 / 5  # TDMA节点个数/总的节点个数
#
#     if env.z < 0:
#         TONT = p * env.z_num1 + (1 - p) * env.z_num3
#     else:
#         TONT = env.z_num1
#
#     return TONT
#
# def my_plot(file1, file2, file3, file4, file5):
#     # 设置最大迭代次数
#     max_iter = 5000
#
#     # 从文本文件中加载智能体、TDMA 和 SA 的奖励数据
#     agent_reward = np.loadtxt(file1)
#     tdma_reward = np.loadtxt(file2)
#     sa_reward = np.loadtxt(file3)
#
#     # 加载已经平滑处理过的 SA 和 TDMA 的奖励数据
#     sa_reward_smoothed = np.loadtxt(file4)
#     tdma_reward_smoothed = np.loadtxt(file5)
#
#     # 计算智能体、TDMA 和 SA 的吞吐量
#     throughput_agent = np.cumsum(agent_reward) / np.arange(1, len(agent_reward) + 1)
#     throughput_tdma = np.cumsum(tdma_reward) / np.arange(1, len(tdma_reward) + 1)
#     throughput_sa = np.cumsum(sa_reward) / np.arange(1, len(sa_reward) + 1)
#
#     # 绘制吞吐量曲线
#     plt.plot(throughput_agent + throughput_tdma + throughput_sa, color='#236B8E', lw=1.2,
#              label='Total')
#     plt.plot(sa_reward_smoothed, color='orange', lw=1.2, label='slottedALOHA')
#     plt.plot(tdma_reward_smoothed, color='green', lw=1.2, label='TDMA')
#
#     # 每隔1000个时隙添加一个原点标记
#     indices = np.arange(0, max_iter, 1000)
#     plt.scatter(indices, throughput_agent[indices] + throughput_tdma[indices] + throughput_sa[indices],
#                 color='#236B8E', s=30, marker='o')  # Total 曲线添加原点标记
#
#     plt.scatter(indices, sa_reward_smoothed[indices], color='orange', s=30, marker='o')  # slottedALOHA 曲线添加原点标记
#     plt.scatter(indices, tdma_reward_smoothed[indices], color='green', s=30, marker='o')  # TDMA 曲线添加原点标记
#
#     # 设置图例、坐标轴和标题
#     plt.legend()
#     plt.xlim((0, max_iter))
#     plt.ylim((0, 1))
#     plt.xlabel('Time Slot')
#     plt.ylabel('Throughput')
#     plt.title('Network Throughput Comparison')
#
#
# if __name__ == '__main__':
#     env1 = ENVIRONMENT()
#     TONT = calculate_TONT(env1)
#     print("TONT =", TONT)
#
#     plt.figure()
#     my_plot('rewards/agent_len5e4_M40.txt',
#             'rewards/TDMA_len5e4_M40.txt',
#             'rewards/SA_len5e4_M40.txt',
#             'rewards/only_SA.txt',
#             'rewards/only_TDMA.txt')
#
#     # 显示图形
#     plt.show()
import numpy as np
import matplotlib.pyplot as plt
from TDMA import run_tdma_simulation
from env import ENVIRONMENT

def calculate_TONT(env):
    p = 2 / 5  # TDMA节点个数/总的节点个数

    if env.z < 0:
        TONT = p * env.z_num1 + (1 - p) * env.z_num3
    else:
        TONT = env.z_num1

    return TONT

def my_plot(file1, file2, file3, file4, file5):
    # 设置最大迭代次数
    max_iter = 10000

    # 从文本文件中加载智能体、TDMA 和 SA 的奖励数据
    agent_reward = np.loadtxt(file1)
    tdma_reward = np.loadtxt(file2)
    sa_reward = np.loadtxt(file3)

    # 加载已经平滑处理过的只有 SA 和 TDMA 的奖励数据
    sa_reward_smoothed = np.loadtxt(file4)
    tdma_reward_smoothed = np.loadtxt(file5)

    # 计算智能体、TDMA 和 SA 的吞吐量
    throughput_agent = np.cumsum(agent_reward) / np.arange(1, len(agent_reward) + 1)
    throughput_tdma = np.cumsum(tdma_reward) / np.arange(1, len(tdma_reward) + 1)
    throughput_sa = np.cumsum(sa_reward) / np.arange(1, len(sa_reward) + 1)

    # 绘制吞吐量曲线
    total_line, = plt.plot(np.arange(1, max_iter + 1), throughput_agent + throughput_tdma + throughput_sa,
                           color='#236B8E', lw=1.2,label='Total Throughput')
    sa_line, = plt.plot(np.arange(1, max_iter + 1), sa_reward_smoothed, color='orange', lw=1.2, label='Slotted ALOHA')
    tdma_line, = plt.plot(np.arange(1, max_iter + 1), tdma_reward_smoothed, color='green', lw=1.2, label='TDMA')


    # 设置图例、坐标轴和标题
    plt.legend(handles=[total_line, sa_line, tdma_line])
    plt.xlim((0, max_iter))
    plt.ylim((0, 1))
    plt.xlabel('Time Slot')
    plt.ylabel('Throughput')
    plt.title('Network Throughput Comparison')

    # 显示图例
    plt.legend()


env1 = ENVIRONMENT()
TONT = calculate_TONT(env1)
print("TONT =", TONT)

plt.figure()
my_plot('rewards/agent_len5e4_M40.txt',
            'rewards/TDMA_len5e4_M40.txt',
            'rewards/SA_len5e4_M40.txt',
            'rewards/only_SA.txt',
            'rewards/only_TDMA.txt')

    # 显示图形
plt.show()
